import type { Reducer } from 'redux';
import { isAction } from './util';
import type { DataModelingThunkAction } from './reducer';
import { analyzeDocuments, type MongoDBJSONSchema } from 'mongodb-schema';
import { getCurrentDiagramFromState } from './diagram';
import { UUID } from 'bson';
import {
  DEFAULT_IS_EXPANDED,
  type Relationship,
} from '../services/data-model-storage';
import { applyLayout } from '@mongodb-js/compass-components';
import {
  collectionToBaseNodeForLayout,
  relationshipToDiagramEdge,
} from '../utils/nodes-and-edges';
import { inferForeignToLocalRelationshipsForCollection } from './relationships';
import { mongoLogId } from '@mongodb-js/compass-logging/provider';

export type AnalysisProcessState = {
  currentAnalysisOptions:
    | ({
        name: string;
        connectionId: string;
        database: string;
        collections: string[];
      } & AnalysisOptions)
    | null;
  analysisProcessStatus: 'idle' | 'in-progress';
  samplesFetched: number;
  schemasAnalyzed: number;
  relationsInferred: number;
};

export const AnalysisProcessActionTypes = {
  ANALYZING_COLLECTIONS_START:
    'data-modeling/analysis-stats/ANALYZING_COLLECTIONS_START',
  NAMESPACE_SAMPLE_FETCHED:
    'data-modeling/analysis-stats/NAMESPACE_SAMPLE_FETCHED',
  NAMESPACE_SCHEMA_ANALYZED:
    'data-modeling/analysis-stats/NAMESPACE_SCHEMA_ANALYZED',
  NAMESPACES_RELATIONS_INFERRED:
    'data-modeling/analysis-stats/NAMESPACES_RELATIONS_INFERRED',
  ANALYSIS_FINISHED: 'data-modeling/analysis-stats/ANALYSIS_FINISHED',
  ANALYSIS_FAILED: 'data-modeling/analysis-stats/ANALYSIS_FAILED',
  ANALYSIS_CANCELED: 'data-modeling/analysis-stats/ANALYSIS_CANCELED',
} as const;

export type AnalysisOptions = {
  automaticallyInferRelations: boolean;
};

export type AnalyzingCollectionsStartAction = {
  type: typeof AnalysisProcessActionTypes.ANALYZING_COLLECTIONS_START;
  name: string;
  connectionId: string;
  database: string;
  collections: string[];
  options: AnalysisOptions;
};

export type NamespaceSampleFetchedAction = {
  type: typeof AnalysisProcessActionTypes.NAMESPACE_SAMPLE_FETCHED;
  namespace: string;
};

export type NamespaceSchemaAnalyzedAction = {
  type: typeof AnalysisProcessActionTypes.NAMESPACE_SCHEMA_ANALYZED;
  namespace: string;
};

export type NamespacesRelationsInferredAction = {
  type: typeof AnalysisProcessActionTypes.NAMESPACES_RELATIONS_INFERRED;
  namespace: string;
  count: number;
};

export type AnalysisFinishedAction = {
  type: typeof AnalysisProcessActionTypes.ANALYSIS_FINISHED;
  name: string;
  connectionId: string;
  database: string;
  collections: {
    ns: string;
    schema: MongoDBJSONSchema;
    position: { x: number; y: number };
    isExpanded: boolean;
  }[];
  relations: Relationship[];
};

export type AnalysisFailedAction = {
  type: typeof AnalysisProcessActionTypes.ANALYSIS_FAILED;
  error: Error;
};

export type AnalysisCanceledAction = {
  type: typeof AnalysisProcessActionTypes.ANALYSIS_CANCELED;
};

export type AnalysisProgressActions =
  | AnalyzingCollectionsStartAction
  | NamespaceSampleFetchedAction
  | NamespaceSchemaAnalyzedAction
  | NamespacesRelationsInferredAction
  | AnalysisFinishedAction
  | AnalysisFailedAction
  | AnalysisCanceledAction;

const INITIAL_STATE = {
  currentAnalysisOptions: null,
  analysisProcessStatus: 'idle' as const,
  samplesFetched: 0,
  schemasAnalyzed: 0,
  relationsInferred: 0,
};

export const analysisProcessReducer: Reducer<AnalysisProcessState> = (
  state = INITIAL_STATE,
  action
) => {
  if (
    isAction(action, AnalysisProcessActionTypes.ANALYZING_COLLECTIONS_START)
  ) {
    return {
      ...INITIAL_STATE,
      analysisProcessStatus: 'in-progress',
      currentAnalysisOptions: {
        name: action.name,
        connectionId: action.connectionId,
        database: action.database,
        collections: action.collections,
        automaticallyInferRelations: action.options.automaticallyInferRelations,
      },
    };
  }
  if (isAction(action, AnalysisProcessActionTypes.NAMESPACE_SAMPLE_FETCHED)) {
    return {
      ...state,
      samplesFetched: state.samplesFetched + 1,
    };
  }
  if (isAction(action, AnalysisProcessActionTypes.NAMESPACE_SCHEMA_ANALYZED)) {
    return {
      ...state,
      schemasAnalyzed: state.schemasAnalyzed + 1,
    };
  }
  if (
    isAction(action, AnalysisProcessActionTypes.ANALYSIS_CANCELED) ||
    isAction(action, AnalysisProcessActionTypes.ANALYSIS_FAILED) ||
    isAction(action, AnalysisProcessActionTypes.ANALYSIS_FINISHED)
  ) {
    return {
      ...state,
      analysisProcessStatus: 'idle',
    };
  }
  return state;
};

async function getInitialLayout({
  collections,
  relations,
}: {
  collections: { ns: string; schema: MongoDBJSONSchema; isExpanded: boolean }[];
  relations: Relationship[];
}) {
  const hasRelations = relations.length > 0;
  const nodes = collections.map((coll) => {
    return collectionToBaseNodeForLayout({
      ns: coll.ns,
      jsonSchema: coll.schema,
      displayPosition: [0, 0],
      isExpanded: coll.isExpanded,
    });
  });
  return await applyLayout({
    nodes,
    edges: relations.map((rel) => relationshipToDiagramEdge(rel, false, [])), // nodes are not important here
    direction: hasRelations ? 'STAR' : 'RECTANGLE',
  });
}

export function startAnalysis(
  name: string,
  connectionId: string,
  database: string,
  collections: string[],
  options: AnalysisOptions
): DataModelingThunkAction<
  Promise<void>,
  | AnalyzingCollectionsStartAction
  | NamespaceSampleFetchedAction
  | NamespaceSchemaAnalyzedAction
  | NamespacesRelationsInferredAction
  | AnalysisFinishedAction
  | AnalysisCanceledAction
  | AnalysisFailedAction
> {
  return async (
    dispatch,
    getState,
    {
      connections,
      cancelAnalysisControllerRef,
      logger,
      track,
      dataModelStorage,
      preferences,
    }
  ) => {
    // Analysis is in progress, don't start a new one unless user canceled it
    if (cancelAnalysisControllerRef.current) {
      return;
    }
    const namespaces = collections.map((collName) => {
      return `${database}.${collName}`;
    });
    const cancelController = (cancelAnalysisControllerRef.current =
      new AbortController());
    dispatch({
      type: AnalysisProcessActionTypes.ANALYZING_COLLECTIONS_START,
      name,
      connectionId,
      database,
      collections,
      options,
    });
    try {
      let relations: Relationship[] = [];
      const dataService = connections.getDataServiceForConnection(connectionId);

      const collections = await Promise.all(
        namespaces.map(async (ns) => {
          const sample = await dataService.sample(
            ns,
            { size: 100 },
            { promoteValues: false },
            {
              abortSignal: cancelController.signal,
              fallbackReadPreference: 'secondaryPreferred',
            }
          );

          const accessor = await analyzeDocuments(sample, {
            signal: cancelController.signal,
          });

          // TODO(COMPASS-9314): Update how we show analysis progress.
          dispatch({
            type: AnalysisProcessActionTypes.NAMESPACE_SAMPLE_FETCHED,
            namespace: ns,
          });

          const schema = await accessor.getMongoDBJsonSchema({
            signal: cancelController.signal,
          });
          dispatch({
            type: AnalysisProcessActionTypes.NAMESPACE_SCHEMA_ANALYZED,
            namespace: ns,
          });
          return { ns, schema, sample, isExpanded: DEFAULT_IS_EXPANDED };
        })
      );

      if (
        preferences.getPreferences().enableAutomaticRelationshipInference &&
        options.automaticallyInferRelations
      ) {
        relations = (
          await Promise.all(
            collections.map(
              async ({
                ns,
                schema,
                sample,
              }): Promise<Relationship['relationship'][]> => {
                const relationships =
                  await inferForeignToLocalRelationshipsForCollection(
                    ns,
                    schema,
                    sample,
                    collections,
                    dataService,
                    cancelController.signal,
                    (err) => {
                      logger.log.warn(
                        mongoLogId(1_001_000_371),
                        'DataModeling',
                        'Failed to identify relationship for collection',
                        { ns, error: err.message }
                      );
                    }
                  );
                dispatch({
                  type: AnalysisProcessActionTypes.NAMESPACES_RELATIONS_INFERRED,
                  namespace: ns,
                  count: relationships.length,
                });
                return relationships;
              }
            )
          )
        ).flatMap((relationships) => {
          return relationships.map((relationship) => {
            return {
              id: new UUID().toHexString(),
              relationship,
              isInferred: true,
            };
          });
        });
      }

      if (cancelController.signal.aborted) {
        throw cancelController.signal.reason;
      }

      const positioned = await getInitialLayout({
        collections,
        relations,
      });

      dispatch({
        type: AnalysisProcessActionTypes.ANALYSIS_FINISHED,
        name,
        connectionId,
        database,
        collections: collections.map((coll) => {
          const node = positioned.nodes.find((node) => {
            return node.id === coll.ns;
          });
          const position = node ? node.position : { x: 0, y: 0 };
          return { ...coll, position };
        }),
        relations,
      });

      track('Data Modeling Diagram Created', {
        num_collections: collections.length,
      });

      void dataModelStorage.save(getCurrentDiagramFromState(getState()));
    } catch (err) {
      if (cancelController.signal.aborted) {
        dispatch({ type: AnalysisProcessActionTypes.ANALYSIS_CANCELED });
      } else {
        logger.log.error(
          mongoLogId(1_001_000_350),
          'DataModeling',
          'Failed to analyze schema',
          { err }
        );
        dispatch({
          type: AnalysisProcessActionTypes.ANALYSIS_FAILED,
          error: err as Error,
        });
      }
    } finally {
      cancelAnalysisControllerRef.current = null;
    }
  };
}

export function retryAnalysis(): DataModelingThunkAction<void, never> {
  return (dispatch, getState) => {
    const { currentAnalysisOptions } = getState().analysisProgress;
    if (!currentAnalysisOptions) {
      return null;
    }
    const {
      name,
      connectionId,
      database,
      collections,
      automaticallyInferRelations,
    } = currentAnalysisOptions;
    void dispatch(
      startAnalysis(name, connectionId, database, collections, {
        automaticallyInferRelations,
      })
    );
  };
}

export function cancelAnalysis(): DataModelingThunkAction<void, never> {
  return (_dispatch, _getState, { cancelAnalysisControllerRef }) => {
    cancelAnalysisControllerRef.current?.abort();
    cancelAnalysisControllerRef.current = null;
  };
}
